import pandas as pd
import numpy as np
import time

class square_env: #环境类，方正格子，内含陷阱和奖励，陷阱结束游戏，奖励也结束游戏
    def __init__(self,trap_loc=None,treasure_loc=None,agent_loc=None,height=10,width=10):
        self.height=height #tuple
        self.width=width
        self.trapLoc=trap_loc #[tuples]
        self.treasureLoc=treasure_loc #[tuples]
        self.agentLoc=agent_loc #[a,b],看作环境的一部分，而外部操作只输入action

    def showEnv(self): #打印环境，左上角为坐标(0,0)
        table=np.zeros((self.height,self.width))
        for i in self.trapLoc:
            table[i]=-1
        for i in self.treasureLoc:
            table[i]=1
        table[tuple(self.agentLoc)]=2
        if tuple(self.agentLoc) in self.trapLoc: #踩到了陷阱
            table[tuple(self.agentLoc)]=-2
        elif tuple(self.agentLoc) in self.treasureLoc: #找到了宝藏
            table[tuple(self.agentLoc)]=3
        else:
            pass
        for i in table:
            str=""
            for j in i:
                if j==0:
                    str+=" _ "
                elif j==-1:
                    str+=" # "
                elif j==1:
                    str+=" T "
                elif j==2:
                    str+=" * "
                elif j==-2:
                    str+=" X "
                elif j==3:
                    str+=" $ "
                else:
                    pass
            print(str)

    def actionResp(self,action): #环境针对外部输入action调整agent的位置并给出响应
        Over=True
        if action == 'up':
            if self.agentLoc[0]==0:
                pass
            else:
                self.agentLoc[0]-=1
        elif action == 'left':
            if self.agentLoc[1]==0:
                pass
            else:
                self.agentLoc[1]-=1
        elif action == 'down':
            if self.agentLoc[0]==self.height-1:
                pass
            else:
                self.agentLoc[0]+=1
        elif action == 'right':
            if self.agentLoc[1]==self.width-1:
                pass
            else:
                self.agentLoc[1]+=1
        else: #未来也许还有停顿这个动作
            pass
        if tuple(self.agentLoc) in self.trapLoc:
            return Over , -10
        elif tuple(self.agentLoc) in self.treasureLoc:
            return Over,10
        else:
            return not Over,-1 #负反馈非常重要，如果每一步没有负反馈（reward=0），Sarsa Lamda的收敛将十分困难，会遇到学习陷阱，跳不出来

    def getAgentLoc(self): #反馈环境上agent的最新位置
        return self.agentLoc

    def reSetEnv(self,agent_loc):
        self.agentLoc=agent_loc

        class agent:
            actions=('up', 'left', 'down', 'right')
            epsilon=.1 #学习时有epsilon的概率使用随机步骤
            lr=0.9 #学习率
            lamda=0.9 #折现率
            sarsaLamdaparam=0.9
            def __init__(self,loction=None):
                self.location=loction[:] #记录自己的位置，因为list在复制时都是引用，所以这里用[:]的方式来浅拷贝
                self.vTable={} #dictionary，价值表的内容都是针对下一步action后得到的价值折现评估
                self.locationNextMove=self.chooseAction() #这个变量是为Sarsa算法服务的，用于记录下一状态的确定了的行动
                self.sarsaLamda={} #dictionary，历史状态记录表，每次学习回合都要恢复,数据结构应该和vTable是一样的
                self.sarsaLamda[tuple(self.location)]=[0,0,0,0]
            def chooseAction(self,how=None):
                if self.vTable.get(tuple(self.location))==None:
                    self.vTable[tuple(self.location)]=[0,0,0,0]
                CurrentLocationValueTable=self.vTable[tuple(self.location)]
                #'''
                if np.random.rand()<=self.epsilon :
                    action=np.random.choice(self.actions)
                else:
                    action=np.random.choice(self.findMaxAction(CurrentLocationValueTable))
                #'''
                #action=np.random.choice(self.findMaxAction(CurrentLocationValueTable)) #如果不使用随机方式，会收敛到局部值
                return action

            def dl(self,action,location,newLocation,how,env,isOver,reward):
                if self.vTable.get(tuple(location))==None:
                    self.vTable[tuple(location)]=[0,0,0,0]
                if self.vTable.get(tuple(newLocation))==None:
                    self.vTable[tuple(newLocation)]=[0,0,0,0]
                if how=="Q":
                    '''
                    Q-learning，缺点是更新慢，每次只能更新最近一个state，每次要更新一个已知value最近的state，
                    需要临近的state根据策略选中这个state，如果临近state可选下一个state很多的话，选中已知state的概率很低，
                    所以开始的学习速度很慢，进化的sarsa lamda或者Q lamda算法会记录历史过去，一起更新，这样收敛比较快
                    '''
                    q_predict=self.vTable[tuple(location)][self.actions.index(action)] #采用已经知道的值来作为最佳估计
                    if isOver != True:
                        q_real=reward+self.lamda*np.array(self.vTable[tuple(newLocation)]).max()
                    else:
                        q_real=reward
                    #table[S,a]+=误差delta
                    self.vTable[tuple(location)][self.actions.index(action)]+=self.lr*(q_real-q_predict)
                elif how=="S": #Sarsa
                    '''
                    通过实验，Q-learning学习时偏重于exploration（深入挖掘），Sarsa偏重于exploration（横向拓展）。
                    在找到最优解后，Q-learning很难学习到新的东西，Sarsa则还会走随机步，去学习新的东西。
                    Q-learning学习时收敛的更快，在找到最优解后，每次学习的偏差不大。
                    Sarsa则依然偏差很大，所以有时候不知道Sarsa是否已经找到了最优解。
                    '''
                    q_predict=self.vTable[tuple(location)][self.actions.index(action)]
                    if isOver != True:
                        q_real=reward+self.lamda*self.vTable[tuple(newLocation)][self.actions.index(self.locationNextMove)]
                    else:
                        q_real=reward
                    self.vTable[tuple(location)][self.actions.index(action)]+=self.lr*(q_real-q_predict)
                elif how=="SarsaLamda":
                    q_predict=self.vTable[tuple(location)][self.actions.index(action)]
                    if isOver != True:
                        q_real=reward+self.lamda*self.vTable[tuple(newLocation)][self.actions.index(self.locationNextMove)]
                    else:
                        q_real=reward
                    delt=q_real-q_predict
                    self.updateSarsaLamda(location,action)
                    for i in self.sarsaLamda:
                        self.vTable[i]=(np.array(self.vTable[i])+self.lr*delt*np.array(self.sarsaLamda[i])).tolist()
                        self.sarsaLamda[i]=(self.lamda*self.sarsaLamdaparam*np.array(self.sarsaLamda[i])).tolist()
                else:
                    pass

            def moveAndLearn(self,how,env):
                if how=="Q": #Q-learning,预估的newLocation价值依赖的那步并非在newlocation时真正会走的那步
                    action=self.chooseAction()
                    beforeAction=self.location[:]
                    isOver,reward=env.actionResp(action)
                    afterAction=env.getAgentLoc()[:]
                    self.dl(action,beforeAction,afterAction,how,env,isOver,reward)
                    self.location=afterAction
                    return isOver
                elif how=="S" or how=="SarsaLamda": #Sarsa,预估的newLocation价值依赖的那步就是未来真正会走的那步（包含随机步的价值）
                    '''
                    Q-learning不能使用批量更新，因为学习是off-policy的，policy仅仅是评估时参考使用，无法将计算得到的价值进行反向追溯。
                    而Sarsa是on-policy，policy得出什么结论，下一步就做什么，最终价值可以反向追溯。
                    Sarsa Lamda方法使用批量更新加快了Sarsa更新速度，是的Sarsa更稳定
                    '''
                    action=self.locationNextMove
                    beforeAction=self.location[:]
                    isOver,reward=env.actionResp(action)
                    afterAction=env.getAgentLoc()[:]
                    self.location=afterAction[:]
                    self.locationNextMove=self.chooseAction() #先记录afterAction这个位置会采用的action，dl里会用到
                    self.dl(action,beforeAction,afterAction,how,env,isOver,reward)
                    return isOver
                else:
                    pass

            def setLocation(self,location):
                self.location=location

            def findMaxAction(self,tlist): #根据vTable里的location的值，将对应的actions中的最大值的动作全部返回
                imax=max(tlist)
                actions_=[]
                count=0
                for i in tlist:
                    if i==imax:
                        actions_.append(self.actions[count])
                    count+=1
                return actions_

            def resetSarsaLamda(self):
                self.sarsaLamda={}

            def updateSarsaLamda(self,location,action): #更新算法的追溯表
                if self.sarsaLamda.get(tuple(location))==None:
                    self.sarsaLamda[tuple(location)]=[0,0,0,0] #使用numpy更方便一些
                index=self.actions.index(action)
                self.sarsaLamda[tuple(location)][index]=1 #多次进入同一状态视作reset，而不是叠加
                #self.sarsaLamda[tuple(location)][index]+=1 #叠加和不叠加应该根据实际情况来确定


episodes=500
start_location=[0,0]
trap_loc=[(3,2),(2,3)]
treasure_loc=[(3,3)]
#how="Q"
how="S"
#how="SarsaLamda"
#总体来说Q速度>SarsaLamda>S
agentO=agent(start_location[:])
steps=[]
for episode in range(episodes):
    env=square_env(trap_loc,treasure_loc,start_location[:],5,5) #环境其实也可以增加一个reset方法，就不用每次新实例化一个出来
    agentO.setLocation(start_location[:])
    over=False
    print(chr(27) + "[2J")
    print("episode:",episode)
    env.showEnv()
    #time.sleep(.3)
    i=0
    while not over:
        i+=1
        over=agentO.moveAndLearn(how,env)
        print(chr(27) + "[2J")
        print("episode:",episode,"  step:",i)
        env.showEnv()
        #time.sleep(.3)
    if how=="SarsaLamda":
        agentO.resetSarsaLamda()
    steps.append(i) #方便看一下每一回合到底用了多少步
